# -*- coding: utf-8 -*-
__all__ = ['kron_prod', 'Smolyak_sparse_grids']

from typing import Sequence
import math

import numpy as np
from numpy.typing import NDArray

def kron_prod(level_seq: Sequence[int], sym: bool = False, quad_type: str = 'uniform') -> tuple[NDArray, NDArray]:
	match quad_type.lower():
		case 'legendre' | 'uniform':
			quad_func = Gaussian_Legendre
		case _:
			quad_func = Gaussian_Hermite
	nodes, weights = quad_func(level_seq[0], sym=sym)
	nodes = np.expand_dims(nodes, axis=0)
	for i in range(1, len(level_seq)):
		next_node, next_weight = quad_func(level_seq[i], sym=sym)
		nodes = np.vstack((np.kron(nodes, np.ones_like(next_node)),
						   np.kron(np.ones(nodes.shape[1]), next_node)))
		weights = np.kron(weights, next_weight)
	return nodes.T, weights

def multi_index(d: int, n: int) -> NDArray:
	"""
	生成所有的和大于等于n, n - d + 1的长度为dim多重下标, 即(i1, i2, ···, id), 总共有C_{n-1}^{d-1}, 其中i1 + i2 +···+ id = |i|.
    generating multi-index (i1, i2, ···, id) where |i| = i1 + i2 + ··· + id satisfies n ≤ |i| ≤ n-d+1
	>>> np.testing.assert_allclose(multi_index(3, 5), np.array([[3, 1, 1], [2, 2, 1], [2, 1, 2],\
											                    [1, 3, 1], [1, 2, 2], [1, 1, 3]]))
	:param d: 维度; dimension
	:param n: 下标之和; the sum of multi index, that is |i|
	:return: 多重下标(i1, i2, ···, id)
	"""
	max_val = n - d + 1
	seq = np.ones(d, dtype=int)
	seq[0] = max_val
	res = np.expand_dims(seq.copy(), axis=0)
	ptr = 0
	while seq[d - 1] < max_val:
		if ptr == d - 1:
			for i in range(d - 2, -1, -1):
				ptr = i
				if seq[i] != 1: break
			seq[ptr] -= 1
			ptr += 1
			seq[ptr] = max_val - sum(seq[:ptr]) + ptr
			seq[ptr + 1:] = 1
		else:
			seq[ptr] -= 1
			ptr += 1
			seq[ptr] += 1
		res = np.vstack((res, seq))
	return res

def Gaussian_Legendre(level: int, sym: bool = False) -> tuple[NDArray, NDArray]:
	match level:
		case 1:
			abscissas = [0.00000000]
			weights = [1.00000000e+0]
		case 2:
			abscissas = [2.88675135e-1]
			weights = [5.00000000e-1]
		case 3:
			abscissas = [0.00000000e-1, 3.87298335e-1]
			weights = [4.44444444e-1, 2.77777778e-1]
		case 4:
			abscissas = [1.69990522e-1, 4.30568156e-1]
			weights = [3.26072577e-1, 1.73927423e-1]
		case 5:
			abscissas = [0.00000000e-1, 2.69234655e-1, 4.53089923e-1]
			weights = [2.84444444e-1, 2.39314335e-1, 1.18463443e-1]
		case 6:
			abscissas = [1.19309593e-1, 3.30604693e-1, 4.66234757e-1]
			weights = [2.33956967e-1, 1.80380787e-1, 8.56622462e-2]
		case 7:
			abscissas = [0.00000000e-1, 2.02922576e-1, 3.70765593e-1, 4.74553956e-1]
			weights = [2.08979592e-1, 1.90915025e-1, 1.39852696e-1, 6.47424831e-2]
		case 8:
			abscissas = [0.91717321e-1, 2.62766205e-1, 3.98333239e-1, 4.80144928e-1]
			weights = [1.81341892e-1, 1.56853323e-1, 1.11190517e-1, 5.06142681e-2]
		case 9:
			abscissas = [0.00000000e-1, 1.62126712e-1, 3.06685716e-1, 4.18015554e-1, 4.84080120e-1]
			weights = [1.65119678e-1, 1.56173539e-1, 1.30305348e-1, 9.03240803e-2, 4.06371942e-2]
		case 10:
			abscissas = [0.74437169e-1, 2.16697697e-1, 3.39704784e-1, 4.32531683e-1, 4.86953264e-1]
			weights = [1.47762112e-1, 1.34633360e-1, 1.09543181e-1, 7.47256746e-2, 3.33356722e-2]
		case 11:
			abscissas = [0.00000000e-1, 1.34771578e-1, 2.59548065e-1, 3.65076003e-1, 4.43531300e-1, 4.89114329e-1]
			weights = [1.36462543e-1, 1.31402272e-1, 1.16596882e-1, 9.31451055e-2, 6.27901847e-2, 2.78342836e-2]
		case 12:
			abscissas = [0.62616704e-1, 1.83915749e-1, 2.93658977e-1, 3.84951337e-1, 4.52058628e-1, 4.90780317e-1]
			weights = [1.24573523e-1, 1.16746268e-1, 1.01583713e-1, 8.00391643e-2, 5.34696630e-2, 2.35876682e-2]
		case _:
			raise NotImplementedError("maximum supported value of 'level' is 12")
	abscissas, weights = np.asarray(abscissas), np.asarray(weights)
	if sym:
		if level % 2 == 1:
			abscissas = np.insert(abscissas, 0, -abscissas[:0:-1])
			weights = np.insert(weights, 0, weights[:0:-1])
		else:
			abscissas = np.insert(abscissas, 0, -abscissas[::-1])
			weights = np.insert(weights, 0, weights[::-1])
	return abscissas, weights

def Gaussian_Hermite(level: int, sym: bool = False) -> tuple[NDArray, NDArray]:
	match level:
		case 1:
			abscissas = [0.00000000e+00]
			weights = [1.00000000e+00]
		case 2:
			abscissas = [1.00000000e+00]
			weights = [5.00000000e-01]
		case 3:
			abscissas = [0.00000000e+00, 1.73205081e+00]
			weights = [6.66666667e-01, 1.66666667e-01]
		case 4:
			abscissas = [7.41963784e-01, 2.33441422e+00]
			weights = [4.54124145e-01, 4.58758548e-02]
		case 5:
			abscissas = [0.00000000e+00, 1.35562618e+00, 2.85697001e+00]
			weights = [5.33333333e-01, 2.22075922e-01, 1.12574113e-02]
		case 6:
			abscissas = [6.16706590e-01, 1.88917588e+00, 3.32425743e+00]
			weights = [4.08828470e-01, 8.86157460e-02, 2.55578440e-03]
		case 7:
			abscissas = [0.00000000e+00, 1.15440539e+00, 2.36675941e+00, 3.75043972e+00]
			weights = [4.57142857e-01, 2.40123179e-01, 3.07571240e-02, 5.48268856e-04]
		case 8:
			abscissas = [5.39079811e-01, 1.63651904e+00, 2.80248586e+00, 4.14454719e+00]
			weights = [3.73012258e-01, 1.17239908e-01, 9.63522012e-03, 1.12614538e-04]
		case 9:
			abscissas = [0.00000000e+00, 1.02325566e+00, 2.07684798e+00, 3.20542900e+00, 4.51274586e+00]
			weights = [4.06349206e-01, 2.44097503e-01, 4.99164068e-02, 2.78914132e-03, 2.23458440e-05]
		case 10:
			abscissas = [4.84935708e-01, 1.46598909e+00, 2.48432584e+00, 3.58182348e+00, 4.85946283e+00]
			weights = [3.44642335e-01, 1.35483703e-01, 1.91115805e-02, 7.58070934e-04, 4.31065263e-06]
		case 11:
			abscissas = [0.00000000e+00, 9.28868997e-01, 1.87603502e+00, 2.86512316e+00, 3.93616661e+00, 5.18800122e+00]
			weights = [3.69408369e-01, 2.42240300e-01, 6.61387461e-02, 6.72028524e-03, 1.95671930e-04, 8.12184979e-07]
		case 12:
			abscissas = [4.44403002e-01, 1.34037520e+00, 2.25946445e+00, 3.22370983e+00, 4.27182585e+00, 5.50090170e+00]
			weights = [3.21664362e-01, 1.46967048e-01, 2.91166879e-02, 2.20338069e-03, 4.83718492e-05, 1.49992717e-07]
		case _: raise NotImplementedError("maximum supported value of 'level' is 12")
	abscissas, weights = np.asarray(abscissas), np.asarray(weights)
	if sym:
		if level % 2 == 1:
			abscissas = np.insert(abscissas, 0, -abscissas[:0:-1])
			weights = np.insert(weights, 0, weights[:0:-1])
		else:
			abscissas = np.insert(abscissas, 0, -abscissas[::-1])
			weights = np.insert(weights, 0, weights[::-1])
	return abscissas, weights

def Smolyak_sparse_grids(d: int, K: int, quad_type: str = 'gaussian') -> tuple[NDArray, NDArray]:
	"""
	A(K,d)φ:=Σ (-1)^{K+d-1-|i|}·C(d-1, |i|-K)·(Q_i₁⊗···⊗Q_{i_d})φ, K≤|i|≤K+d-1
    生成数值求积的Smolyak稀疏网格的求积节点和对应的权重
	spawns the nodes and weights of Smolyak sparse grid of tensor product quadrature
	:param quad_type:
	:param d: 求积维度; the dimension of quadrature
	:param K: 数值精度, 关系为|k|=2K-1; the sparse level 'K' of degree of accuracy |k| of numerical integration, where |k|=2K-1
	:return: 高斯求积节点和对应的权重; nodes and weights
	"""
	min_imd, max_imd = max(0, K - d), K - 1
	nodes = np.zeros((1, d))
	weights = np.array([0.])
	for i in range(min_imd, max_imd + 1):
		factor = 1 if (max_imd - i) % 2 == 0 else -1  # (-1) ** (max_imd - i)
		factor *= math.comb(d - 1, i + d - K)
		indexes = multi_index(d, i + d)
		for index in indexes:
			node, weight = kron_prod(index, quad_type=quad_type)
			nodes = np.vstack((nodes, node))
			weights = np.append(weights, weight * factor)
	np.delete(nodes, 0, axis=0)
	np.delete(weights, 0)

	ind = np.lexsort(np.flip(nodes.T, axis=0))
	sort_nodes = nodes[ind]
	sort_weights = weights[ind]

	ind = np.array([0], dtype=int)
	last_ind = 0
	for i in range(1, len(sort_nodes)):
		if np.array_equal(sort_nodes[i], sort_nodes[i - 1]):
			sort_weights[last_ind] += sort_weights[i]
		else:
			last_ind = i
			ind = np.append(ind, i)
	nodes = sort_nodes[ind]
	weights = sort_weights[ind]

	num_row = len(weights)
	for col in range(d):
		ind = np.zeros(num_row, dtype=int)
		num_new = 0
		for row in range(num_row):
			if not np.isclose(nodes[row, col], 0):
				ind[num_new] = row
				num_new += 1
		if num_new > 0:
			nodes = np.vstack((nodes, nodes[ind[:num_new]].copy()))
			nodes[num_row:num_row + num_new, col] = -nodes[num_row:num_row + num_new, col]
			weights = np.append(weights, weights[ind[:num_new]])
			num_row += num_new

	ind = np.lexsort(np.flip(nodes.T, axis=0))
	nodes = nodes[ind]
	weights = weights[ind]
	weights /= sum(weights)
	return nodes, weights
